#!/usr/bin/env python3

import sys
sys.path.append('../../../tools')
from gmcore_plot import *

parser = argparse.ArgumentParser(description='Plot test result')
parser.add_argument('-i', '--input', help='Input file', default='dcmip31.h0.nc')
parser.add_argument('-t', '--time-step', help='Time step', type=int, default=-1)
parser.add_argument('-l', '--loop-plot', help='Loop plot', action='store_true')
args = parser.parse_args()

fig = plt.figure(figsize=(10, 8))

# Set background potential temperature and zonal velocity.
lat = 0

X = 125
radius = 6.37122e6 / X
omega = 0
u0 = 20
p0 = 1e5
teq = 300
N = 0.01
t0 = g.magnitude**2 / N**2 / cp.magnitude
f = addfile('dcmip31.h0.nc').isel(time=0).sel(lat=lat, method='nearest')
pt0 = f.pt.copy()
ts = t0 + (teq - t0) * np.exp(-u0 * N**2 / (4 * g.magnitude**2) * (u0 + 2 * omega * radius) * (np.cos(2 * lat) - 1))
for k in range(f.lev.size):
	for i in range(f.lon.size):
		pt0[k,i] = ts * (p0 / f.phs[i])**(rd / cp) / (ts / t0 * ((f.ph[k,i] / f.phs[i])**(rd / cp) - 1) + 1)
u0 = f.u.copy()
f.close()

pt_levels = np.linspace(-0.1, 0.1, 11)
u_levels = np.linspace(-0.25, 0.25, 11)
colors = ListedColormap([
	"#7a199d",
	"#002edd",
	"#0042e0",
	"#00985c",
	"#00d46e",
	"#eefbf2",
	"#fffef1",
	"#fffb55",
	"#ffe900",
	"#ffbe00",
	"#ff6e00",
	"#ff0000"
])

def plot(data_path, fig):
	f = addfile(data_path).isel(time=args.time_step).sel(lat=lat, method='nearest')

	time = parse_datetime(f.time)
	fig.suptitle('Time: %s' % time.strftime('%Y-%m-%d %H:%M:%S'), fontsize=8)

	ax = fig.add_subplot(2, 1, 1)
	dpt = f.pt - pt0
	plot_contour_lon(ax, dpt, levels=pt_levels, cmap=colors, invert_yaxis=True, with_grid=False, with_contour=True)
	ax.set_title(f'Max: {dpt.max().values:.2f} Min: {dpt.min().values:.2f}', fontsize=8)

	ax = fig.add_subplot(2, 1, 2)
	du = f.u - u0
	plot_contour_lon(ax, du, levels=u_levels, cmap=colors, invert_yaxis=True, with_grid=False, with_contour=True)
	ax.set_title(f'Max: {du.max().values:.2f} Min: {du.min().values:.2f}', fontsize=8)

	f.close()

if args.loop_plot:
	loop_plot(args.input, fig, plot)
else:
	plot(args.input, fig)
	plt.tight_layout()
	plt.show()
